# world.py
# A possible world, where all the random variables are assigned a single value

# local imports

# pyblog imports
import glob
from glob import GROUPBY_NOTHING, GROUPBY_UNUSED_ARGS, GROUPBY_INDICATOR_FN
from misc.poly import Polynomial
from misc.code import unused_arguments, indicator_function, indicator_bound

#
# a varkey is a tuple consisting of a variable name and its arguments
# for example f(1,2) would have the varkey (f, 1, 2)
#
# convert a varkey to a string
def str_varkey(v):
  if type(v[0]) is tuple:
    return "(" + v[0][0].__name__ + ")" + str(v[1:])
  else:
    return v[0].__name__ + str(v[1:])

# get the function name in the varkey
def fn_varkey(v): return v[0]

# execute the dependency function for a variable --
# for stochastic dependencies this returns a distribution while for
# deterministic dependencies this returns a value
def exec_dep_varkey(v): return glob.var_def[v[0]].dep_fn(*v[1:])

# does this varkey refer to a variable with a stochastic dependency
def is_stochastic_varkey(v): return glob.var_def[v[0]].is_stochastic

# the dependency type of the variable
VARTYPE_DETERMINISTIC, VARTYPE_STOCHASTIC, VARTYPE_GROUPBY_CHILD,\
                       VARTYPE_GROUPBY_PARENT = range(4)
#
# A Variable is an instantiation of a random variable in a world
#
class Variable:
  """
  varkey   - a variable's key is a tuple of function name and its arguments
  value    - the actual value of this variable
  dist     - the distribution of this variable (None for deterministic dep)
  parents  - List of parent variables
  children - List of child variables
  type     - deterministic dep, stochastic dep, groupby child, groupby parent
  is_obs   - is this an observed variables (aka evidence)
  """
  def __init__(self, varkey):
    self.parents = set()
    self.children = set()
    self.varkey = varkey
    self.type = None
    self.is_obs = False

  def deepcopy(self):
    """
    Returns a deep copy of this variable.
    NOTE: the distribution is copied by reference since its not modifiable.
    """
    # create a new variable
    var = Variable(self.varkey)
    # copy all the variable attributes
    var.parents, var.children = set(self.parents), set(self.children)
    var.dist, var.type, var.is_obs = self.dist, self.type, self.is_obs
    # groupby parents don't have a value
    if hasattr(self, "value"):
      var.value = self.value
    # groupby parents might have sufficient stats for their children
    # which should be copied over
    elif hasattr(self, "child_suff_stats"):
      var.child_suff_stats = self.child_suff_stats.copy()
    return var

  def add_child_stats(self, childval):
    if hasattr(self, "child_suff_stats"):
      self.child_suff_stats.add(childval)

  def upd_child_stats(self, oldval, newval):
    if hasattr(self, "child_suff_stats"):
      self.child_suff_stats.remove(oldval)
      self.child_suff_stats.add(newval)
  
  def add_child_var(self, childvar):
    self.children.add(childvar.varkey)
    childvar.parents.add(self.varkey)
    if hasattr(self, "child_suff_stats") and hasattr(childvar, "value"):
      self.child_suff_stats.add(childvar.value)

  def remove_child_var(self, childvar):
    self.children.remove(childvar.varkey)
    childvar.parents.remove(self.varkey)
    if hasattr(self, "child_suff_stats"):
      self.child_suff_stats.remove(childvar.value)
    
  def __str__(self):
    ans = ""
    ans += str_varkey(self.varkey)
    if hasattr(self, "value"):
      ans += " == " + str(self.value)
    if hasattr(self, "dist"):
      ans += " ~ " + str(self.dist)
    if hasattr(self, "child_suff_stats"):
      ans += " child stats: " + str(self.child_suff_stats)
    return ans

# NewVariable exception is raised when creating a new variable is not allowed
class NewVariable(Exception):
  def __init__(self, value):
    self.value = value
  def __str__(self):
    return repr(self.value)

def analyze_groupby(vardef):
  """
  Determines how a variable should be grouped
  """
  assert(vardef.groupby is None)
  # if the user has disabled lifted variables then simply skip them here
  if not glob.PARAM_LIFTED_VARIABLES:
    vardef.groupby = GROUPBY_NOTHING
    return
  # only variables with stochastic dependency are grouped
  if not vardef.is_stochastic:
    vardef.groupby = GROUPBY_NOTHING
    return
  # first, compute the number of free arguments of the dependency function
  num_unused = unused_arguments(vardef.dep_fn)
  if num_unused:
    vardef.groupby = GROUPBY_UNUSED_ARGS
    vardef.used_args = vardef.argcount - num_unused
    return
  # second, look for an indicator function which is also a random variable
  # with a stochastic dependency
  ind_fn = indicator_function(vardef.dep_fn)
  if ind_fn is not None and ind_fn in glob.var_def \
         and glob.var_def[ind_fn].is_stochastic:
    vardef.groupby = GROUPBY_INDICATOR_FN
    vardef.ind_fn = ind_fn

    # extract the list of bound functions if any
    vardef.bound_fns = indicator_bound(vardef.dep_fn, ind_fn)
    # and check if the bound functions don't use their arguments
    if vardef.bound_fns is not None:
      for bndfn in vardef.bound_fns:
        if (bndfn not in glob.var_def) or \
           (unused_arguments(glob.var_def[bndfn].dep_fn) != 1):
          vardef.bound_fns = None
    
    return
  # otherwise, no grouping
  vardef.groupby = GROUPBY_NOTHING
  return
  
class World:
  """
  A world is a self-supporting instantiation of random variables.
  Self-supporting means that for any variable in the world, its parents
  will also be instantiated in the world.
  """
  def __init__(self, rand, trace=False, deterministic=False):
    """
    rand   - random number generator
    trace  - tracing
    """
    self.rand, self.trace, self.deterministic = rand, trace, deterministic
    self.variables = {}                 # all the variables
    self.sampleable_vars = set()        # variables to be sampled
    self.initial_vals = {}              # initial values for variables
    self.last_var = None
    self.raise_on_new_var = False       # variables can be created
    self.last_indicator_fnval = None    # for constructing grouby vars

  def _init_scan_list(self):
    """
    Internal function called by init_scan and next_scan to initialize
    the list of variables in the world.
    """
    self.scan_list = list(self.sampleable_vars)
    # if determinism is requested then sort the sampled variables
    # to remove a potential source of non-determinism
    if self.deterministic:
      self.scan_list.sort(key=lambda x: (x[0].__name__,) + x[1:])
    self.scan_idx = 0                   # next variable to flip

  def init_scan(self):
    """
    Initiate a scan of the variables in the world by first
    calling init_scan. Subsequently next_scan_var can be called
    till None is returned. At that point next_scan will have to be called
    to start a new scan.
    """
    self._init_scan_list()
    self.scan_cnt = 0                   # number of scans
    
  def next_scan(self):
    """
    Call next_scan after next_scan_var returns None
    """
    self._init_scan_list()
    self.scan_cnt +=1                   # number of scans

  def next_scan_var(self):
    """
    Either return the next variable in the scan or None if a
    scan has ended. init_scan should have been called before the first
    call to next_scan_var.
    """
    while self.scan_idx < len(self.scan_list):
      varkey = self.scan_list[self.scan_idx]
      self.scan_idx += 1
      # the variable may have been uninstantiated since the time that the
      # scan list was created, so check if its still around
      if varkey in self.variables:
        return self.variables[varkey]
    else:
      return None
    
  def get_log_likelihood(self):
    """
    Likelihood is defined as the joint probability of the evidence
    given all the latent variables in the model. Note: if there is no
    evidence then the likelihood is 1. Also, if there is evidence on a
    deterministic variable then the likelihood is not defined.
    
    NOTE: If the model contains unnormalized distributions then the log
    likelihood may actually be positive. However, as long as the distribution
    for each evidence variable doesn't actually change (the distribution
    parameters might though) we can use the log likelihood in a relative
    sense.
    """
    total = 0.0
    for var in self.variables.itervalues():
      # look for observed variables
      if not var.is_obs:
        continue
      # deterministic variables are not supported as evidence
      if var.type == VARTYPE_DETERMINISTIC:
        raise TypeError("Log Likelihood can't be computed for evidence var" \
                        + " with deterministic dependency:", str(var))
      # get the distribution
      if type(var.dist) is tuple:
        dist = self.variables[var.dist].dist
      else:
        dist = var.dist
      total += dist.log_prob(var.value)

    return total

  def get_log_probability(self):
    """
    Returns the joint log probability of the world.
    If the model contains distributions which are not normalized then
    the log probability might actually be positive!
    """
    total = 0.0
    for var in self.variables.itervalues():
      if var.dist is None:
        continue               # ignore variables with deterministic dependency
      # for groupby children the distribution is kept in a different variable
      if type(var.dist) is tuple:
        dist = self.variables[var.dist].dist
      else:
        dist = var.dist
      # compute the log probability of this variable's value
      lp = dist.log_prob(var.value)
      # add to the total
      total += lp

    return total

  def get_all_vars(self):
    """
    Returns a hash table with the names of the variables and their values
    """
    return dict((x.varkey, x.value) for x in self.variables.itervalues() \
                if hasattr(x, "value"))
  
  def set_observation(self, obslist):
    for obs in obslist:
      if hasattr(obs, "name") and hasattr(obs, "value"):
        self.observe_var(obs.name(), obs.value())
      else:
        raise TypeError("Observation must be RandomVar(..) == val")

  def set_init_vals(self, init_vals):
    for iv in init_vals:
      if hasattr(iv, "name") and hasattr(iv, "value"):
        self.initial_vals[iv.name()] = iv.value()
      else:
        raise TypeError("Initial value must be RandomVar(..) == val")

  def get_children_rec(self, var):
    """
    Returns all the children of this variable, and for deterministic children,
    recurses down to them.
    """
    childkeys = set()
    queue = [var]
    while queue:
      currvar = queue.pop(0)
      for child in currvar.children:
        if child in childkeys:
          continue
        # add this child to the answer set
        childkeys.add(child)
        # lookup child variable
        childvar = self.variables[child]
        # recurse for deterministic children
        if childvar.type == VARTYPE_DETERMINISTIC:
          queue.append(childvar)
            
    return childkeys

  def get_children_values(self, var):
    """
    Returns the values of all the children of this variable
    """
    childvals = []
    # scan all the children
    for childkey in var.children:
      childvar = self.variables[childkey]
      # for groupby children, scan the grandchildren
      if childvar.type == VARTYPE_GROUPBY_PARENT:
        for grandchildkey in childvar.children:
          grandchildvar = self.variables[grandchildkey]
          childvals.append(grandchildvar.value)
      else:
        childvals.append(childvar.value)
    # finally return the values
    return childvals

  def get_collapsed_dist(self, lifted_var, childkey):
    """
    The collapsed dist of a lifted variable returns the variable's distribution
    after integrating away the value of its parent. Currently this is only
    supported for variables which 
    - have one parent
    - are the only child of the parent
    
    If the parent can't be collapsed then the distribution of this variable
    will be returned
    """
    # check that this lifted variable has only one parent and is the only
    # child of its parent and has sufficient stats
    if not glob.PARAM_COLLAPSED_SAMPLING or (len(lifted_var.parents) != 1) \
       or (not hasattr(lifted_var, "child_suff_stats")):
      return lifted_var.dist

    parentkey, = list(lifted_var.parents)
    parent = self.variables[parentkey]

    # check if its the only child
    if parent.children != set([lifted_var.varkey]):
      return lifted_var.dist

    # now try to change the value of the parent to its distribution
    # this triggers it to return a collapsed distribution if possible
    try:
      saved_parent_val = parent.value
      parent.value = self.get_dist(parent)
      # can it be collapsed?
      try:
        tmp = exec_dep_varkey(childkey)
      except:
        return lifted_var.dist
      collapsed = tmp(lifted_var.child_suff_stats)
      return collapsed
    finally:
      parent.value = saved_parent_val
    
  def get_dist(self, var):
    """
    Returns the distribution for a variable
    """
    if var.type == VARTYPE_GROUPBY_CHILD:
      return self.variables[var.dist].dist
    else:
      return var.dist
                        
  def execute_dep(self, var):
    """
    Execute the dependency statement of the variable. This will
    populate the parent and child relationships between var and its
    supporting ancestors, also var.dist and var.type. This function
    also populates var.value for deterministic dependencies
    """
    # initialize last_var to compute parent child pointers
    self.last_var = var
    
    varkey = var.varkey
    
    # is this a group-by var?
    if type(varkey[0]) is tuple:
      vardef = glob.var_def[varkey[0][0]]

      # set the variable's type and value
      var.type = VARTYPE_GROUPBY_PARENT
      
      if vardef.groupby == GROUPBY_UNUSED_ARGS:
        # construct the arguments by putting None for the unused arguments,
        # they are "unused" right?
        args = varkey[1:] + \
               tuple([None for _ in range(vardef.argcount - vardef.used_args)])
        # now execute the dependency statement to get the distribution
        var.dist = vardef.dep_fn(*args)

      # GROUPBY_INDICATOR_FN
      else:
        # construct the args with None, the arguments will only be
        # passed to the indicator function which will always evaluate
        # to the fixed value because we are setting last_indicator_fnval
        args = [None for _ in range(vardef.argcount)]
        self.last_indicator_fnval = (vardef.ind_fn, varkey[1])
        # execute the dependency statement
        var.dist = vardef.dep_fn(*args)
        self.last_indicator_fnval = None

      # for group-by functions we'll maintain sufficient stats if possible
      if hasattr(var.dist, "SufficientStats"):
        statsclass = var.dist.SufficientStats
        # if there was no existing child sufficient stats or
        # if the existing stats were for a different distribution
        # then we need to construct the stats again from scratch
        if (not hasattr(var, "child_suff_stats")) or \
           (not isinstance(var.child_suff_stats, statsclass)):
          var.child_suff_stats = statsclass()
          for childkey in var.children:
            childvar = self.variables[childkey]
            var.child_suff_stats.add(childvar.value)
      # remove any child sufficient stats if one existed
      elif hasattr(var, "child_suff_stats"):
        del var.child_suff_stats
    
    # non-groupby var
    else:
      vardef = glob.var_def[varkey[0]]
    
      # stochastic dependency
      if vardef.is_stochastic:
        # check if the groupby needs to be analyzed
        if vardef.groupby is None:
          analyze_groupby(vardef)
      
        # plain stochastic variable
        if vardef.groupby == GROUPBY_NOTHING:
          var.dist = vardef.dep_fn(*varkey[1:])
          var.type = VARTYPE_STOCHASTIC

        # grouped stochastic variables

        elif vardef.groupby == GROUPBY_UNUSED_ARGS:
          var.dist = ((varkey[0],),) + varkey[1:vardef.used_args+1]
          # lookup groupby variable and create parent child edges
          self.query_var(var.dist, retval = False)
          # set the variable's type
          var.type = VARTYPE_GROUPBY_CHILD
          
        elif vardef.groupby == GROUPBY_INDICATOR_FN:
          # compute the indicator variable's key
          ind_key = (vardef.ind_fn,) + varkey[1:]
          # lookup the indicator variable (and also create parent, child
          # edges to the indicator variable)
          ind_val = self.query_var(ind_key)
          # now compute the grouby variable's key
          var.dist = ((varkey[0],), ind_val)
          # lookup groupby variable and create parent child edges
          self.query_var(var.dist, retval=False)
          # set the variable's type
          var.type = VARTYPE_GROUPBY_CHILD
      
      # deterministic dependency
      else:
        var.dist = None
        var.value = vardef.dep_fn(*varkey[1:])
        var.type = VARTYPE_DETERMINISTIC
      
    # done computing parent child pointers, clear last_var
    self.last_var = None
  
  def new_var(self, varkey, observed=False):
    """
    Create a new variable, compute its distribution and populate its
    parent pointers (and child ptrs from the parents)
    
    For variables with deterministic dependency it computes the value rather
    than the distribution.
    """
    # check that we are currently allowed to create variables
    if self.raise_on_new_var:
      raise NewVariable, varkey
    
    # create a new variable object
    var = Variable(varkey)
    var.is_obs = observed               # is this an observation?
      
    if self.trace:
      print "new_var:", var

    # execute the dependency statement and populate parent child
    # pointers and either the dist or the value
    self.execute_dep(var)

    # now store the variable in the world
    self.variables[varkey] = var

    # mark it sampleable if appropriate
    if not observed \
       and var.type in (VARTYPE_STOCHASTIC, VARTYPE_GROUPBY_CHILD):
      self.sampleable_vars.add(varkey)
    
    return var
    
  def observe_var(self, varkey, val):
    # check for duplicates
    if varkey in self.variables:
      raise Exception("Duplicate observation of %s" % str(Variable(varkey)))

    # create a new variable and record its value
    var = self.new_var(varkey, True)
    if var.dist is None:
      raise TypeError("Evidence can only be on variables with stochastic"\
                      +" dependency")
    var.value = val
    if var.type == VARTYPE_GROUPBY_CHILD:
      gby_parent = self.variables[var.dist]
      gby_parent.add_child_stats(var.value)

    if self.trace:
      print "observe_var:", var

  def remove_var(self, var):
    if self.trace:
      print "remove_var:", str(var)

    # better not have children or be evidence
    assert(not var.is_obs and not var.children)
    
    del self.variables[var.varkey]
    # mark it as not sampleable if it was marked sampleable earlier
    if var.varkey in self.sampleable_vars:
      self.sampleable_vars.remove(var.varkey)

    # remove from the parent and recursively remove parent
    for parkey in list(var.parents):
      par = self.variables[parkey]
      par.remove_child_var(var)
      # we can only remove the parent if its not an observation
      if not par.is_obs and not par.children:
        self.remove_var(par)

    assert(not len(var.parents))
    # finally delete the variable
    del var

  def query_var(self, varkey, retval=True):
    # if an indicator function's value has been fixed then simply return
    # the fixed value
    if (self.last_indicator_fnval) is not None \
       and (varkey[0] == self.last_indicator_fnval[0]):
      return self.last_indicator_fnval[1]
    
    last_var = self.last_var            # save last_var
    last_indicator_fnval = self.last_indicator_fnval
    self.last_indicator_fnval = None
    
    # does the variable exist?
    if varkey in self.variables:
      var = self.variables[varkey]

    # if not then create a new variable
    else:
      var = self.new_var(varkey)
      # if there is an initial value for this variable then use that
      if varkey in self.initial_vals:
        var.value = self.initial_vals[varkey]
        if var.type == VARTYPE_GROUPBY_CHILD:
          gby_parent = self.variables[var.dist]
          gby_parent.add_child_stats(var.value)
      # otherwise sample a value
      elif var.type == VARTYPE_STOCHASTIC:
        var.value = var.dist.sample(self.rand)
      # for grouped variables we need to find the distribution first
      elif var.type == VARTYPE_GROUPBY_CHILD:
        gby_parent = self.variables[var.dist]
        var.value = gby_parent.dist.sample(self.rand)
        gby_parent.add_child_stats(var.value)

      if self.trace:
        print "query_var:", var


    # build parent and child relationships
    if last_var is not None:
      var.add_child_var(last_var)

    self.last_indicator_fnval = last_indicator_fnval
    self.last_var = last_var            # restore last_var
    
    # finally, return the value
    if retval:
      return var.value
    else:
      return var

  def dump(self):
    for varkey, var in self.variables.iteritems():
      if var.is_obs:
        print "Observed",
      if var.dist is not None:
        print "@var_dist",
      else:
        print "@var",
      print var
      if var.parents:
        print "  Parents:", [str_varkey(x) for x in var.parents]
      if var.children:
        print "  Children:", [str_varkey(x) for x in var.children]

  def update_nonswitching_var(self, var, newval):
    """
    Update the value of non-switching variables. Switching variables should
    use WorldDiff to change the world.

    The following code assumes that since the variable is nonswitching none
    of the parent child relationships in the world will change as a
    consequence of changing its value.
    """
    # update lifted stats in the parent
    if var.type == VARTYPE_GROUPBY_CHILD:
      gby_parent = self.variables[var.dist]
      gby_parent.upd_child_stats(var.value, newval)
      
    var.value = newval
    for childkey in var.children:
      childvar = self.variables[childkey]
      self.execute_dep(childvar)
      # children with deterministic dependencies will require their
      # children to be updated as well
      if childvar.type == VARTYPE_DETERMINISTIC:
        self.update_nonswitching_var(childvar, childvar.value)
  
  def query_and_forget(self, rvars):
    """
    query some random variables and uninstantiate them from the world
    (unless the variable was already in the world)
    """
    # initialize the answer
    answer = []
    
    # initialize the list of variables to be later deleted
    dellist = []

    for rvar in rvars:
      # newly queried variables need to be deleted later
      if rvar.varkey not in self.variables:
        new = True
      else:
        new = False
      answer.append(self.query_var(rvar.varkey))
      if new:
        dellist.append(self.variables[rvar.varkey])

    for var in dellist:
      # if this variable has children then they would also be on the delete
      # list, so we can skip this for now
      if not var.children:
        self.remove_var(var)

    return answer
    
class WorldDiff(World):
  """
  A WorldDiff is an extension of an underlying world object. In a WorldDiff,
  variables may be changed and parent child relationships might be different.
  However, the underlying world object is not modified unless the WorldDiff
  is merged.

  One can call change_var(..) repeatedly and then finally call close()
  to create a consistent world. No more change_var can be called after close().
  Each call to change_var changes not only the variable to be changed but also
  any deterministic dep children of that variable.

  TODO: change var currently doesn't support deleting and reinstantiating
  a variable and removing said variable from vars_new_parents
  """
  def __init__(self, orig_world):
    """
    orig_world - the original world object which will be modified
    """
    # initialize the parent World object keeping consistent with the
    # original world object
    World.__init__(self, orig_world.rand, orig_world.trace,
                   orig_world.deterministic)
    
    self.orig_world = orig_world        # store the original world
    
    # vars_new_parents are those variables which have taken on new parents.
    self.vars_new_parents = set()

    # vars_new_vals are those variables who have been assigned a new value.
    self.vars_new_vals = set()

    self.created_vars = set()           # newly created variables

    self.deleted_vars = set()           # variables deleted from the old world

    self.closed = False                 # has close() been called

    self.raise_on_new_var = False       # variables can be created

  def pullin_var(self, varkey):
    """
    Return an existing variable or copy from the old world
    The variable better exist in the old world if it doesn't exist in this
    world.
    NOTE: proposal function is not pulled in
    """
    if varkey in self.variables:
      var = self.variables[varkey]
    else:
      if self.raise_on_new_var:
        raise NewVariable, varkey
      # pull-in the variable from the old world and make a copy of it
      var = self.orig_world.variables[varkey].deepcopy()
      # store this copy in the new world
      self.variables[varkey] = var
      # maintain sampleable vars
      if varkey in self.orig_world.sampleable_vars:
        self.sampleable_vars.add(varkey)
      
    return var    

  def change_var(self, varkey, newval):
    """
    varkey - name of the variable
    newval - new value to give to this variable
    """
    if self.trace:
      print "change_var:", str_varkey(varkey)
    assert(not self.closed)
    # pullin this variable from the underlying world
    var = self.pullin_var(varkey)
    # we can only change a sampleable variable
    assert(var.varkey in self.sampleable_vars)
    # update lifted stats in the parent
    if var.type == VARTYPE_GROUPBY_CHILD:
      gby_parent = self.pullin_var(var.dist)
      gby_parent.upd_child_stats(var.value, newval)    
    # assign a new value to this variable
    var.value = newval
    # change child-parent pointers for its children
    self.__change_var(var)
    
  def __change_var(self, var):
    """
    Change an existing variable without setting a value for it
    """
    # add this variable to the set of vars with new values
    self.vars_new_vals.add(var.varkey)
    # record the fact that the children of this variable need new parents
    self.vars_new_parents.update(var.children)

    # pull in all the children of this variable
    for childkey in list(var.children):
      childvar = self.pullin_var(childkey)
      var.remove_child_var(childvar)
      # pullin all the parents of this child
      for parkey in list(childvar.parents):
        parvar = self.pullin_var(parkey)
        parvar.remove_child_var(childvar)
      # no parents left, it will be recomputed on close
      assert(not len(childvar.parents))
      # if the child has a deterministic dependency on its parent then
      # the child needs to be changed as well
      if (childvar.dist is None) and (childkey not in self.vars_new_vals):
        # recursively change the children of a deterministic dep child
        self.__change_var(childvar)
    # no children left; these will be recomputed in close()
    assert(not len(var.children))

  def close(self):
    """
    After close() no more changes can be made on this world. The parent and
    child pointers of all variables will be fixed up correctly and new
    variables may be introduced and some variables may be deleted.
    """
    assert(not self.closed)

    # save the old value of the curr_world variable to restore it
    # in case of any exceptions
    old_curr_world = glob.curr_world
    glob.curr_world = self

    try:
      # compute new parents for the needed variables
      for varkey in self.vars_new_parents:
        var = self.variables[varkey]
        self.execute_dep(var)
    
    finally:
      glob.curr_world = old_curr_world

    # delete variables which don't have any children and are not evidence
    # variables
    for var in [x for x in self.variables.values() \
                if ((not x.is_obs) and (not x.children))]:
      self.remove_var(var)
      
    # now the world is closed to further modification
    self.closed = True
      
  def merge(self):
    assert(self.closed)
    # delete deleted vars from set of variables
    for varkey in self.deleted_vars:
      del self.orig_world.variables[varkey]
      if varkey in self.orig_world.sampleable_vars:
        self.orig_world.sampleable_vars.remove(varkey)

    # for any variable in the new world replace the corresponding variable
    # in the old world
    for varkey, var in self.variables.iteritems():
      self.orig_world.variables[varkey] = var

    # add any new variables to the sampleable set
    self.orig_world.sampleable_vars.update(self.sampleable_vars)
      
  def dump(self):
    print "vars_new_vals:", [str_varkey(x) for x in self.vars_new_vals]
    print "created_vars:", [str_varkey(x) for x in self.created_vars]
    print "deleted_vars:", [str_varkey(x) for x in self.deleted_vars]
    for var in self.variables.itervalues():
      print var
      if var.parents:
        print "  parents:", [str_varkey(x) for x in  var.parents]
      if var.children:
        print "  children:", [str_varkey(x) for x in var.children]
    
  def remove_var(self, var):
    if self.trace:
      print "WorldDiff remove_var:", var
    varkey = var.varkey
    # if this was a newly created variable then we can just remove it as
    # a new variable
    if varkey in self.created_vars:
      self.created_vars.remove(varkey)
    # otherwise we need to mark this for deletion in the old world
    else:
      self.deleted_vars.add(varkey)
    
    if varkey in self.vars_new_vals:
      del self.vars_new_vals[varkey]
      
    if varkey in self.vars_new_parents:
      del self.vars_new_parents[varkey]

    # when we remove a variable its parents need to be modified
    for varkey in var.parents:
      self.pullin_var(varkey)
      
    # super-class remove_var
    World.remove_var(self, var)

  def new_var(self, varkey):
    self.created_vars.add(varkey)
    # super-class new_var method
    return World.new_var(self, varkey)

  def query_var(self, varkey, retval=True):
    # pullin the variable from the original world if at all possible
    try:
      self.pullin_var(varkey)
    except KeyError:
      pass
    
    # super-class query_var
    return World.query_var(self, varkey, retval)
